import optuna
from debias import run
import time
import os
from pathlib import Path
import contextlib
from multiprocessing import Pool

my_dir = Path(__file__).parent
project_dir = my_dir
search_output_dir = project_dir / "optuna-search"
if not search_output_dir.exists():
    os.mkdir(search_output_dir)


class Config:
    pass


def objective(trial, model, dataset, device="auto"):
    cfg = Config()
    batch_size = trial.suggest_int("batch_size", 32, 128)
    learning_rate = trial.suggest_loguniform("learning_rate", 1e-5, 1e-2)
    w_learning_rate = trial.suggest_loguniform("w_learning_rate", 1e-5, 1e-2)

    setattr(cfg, "epoch", 200)
    setattr(cfg, "batch_size", batch_size)
    setattr(cfg, "lr", learning_rate)
    setattr(cfg, "w_lr", w_learning_rate)
    setattr(cfg, "model", model)
    setattr(cfg, "dataset", dataset)
    setattr(cfg, "early_stop", True)
    setattr(cfg, "save", False)
    setattr(cfg, "less_data", False)
    setattr(cfg, "device", device)
    setattr(cfg, "seed", int(time.time()))
    setattr(cfg, "epoch_w", 10)
    setattr(cfg, "load", None)
    result = run(cfg)
    return result[f"{model} on {dataset}"]["auc"]


def run_wrapper(
    study_name, storage, stdout_file, stderr_file, model, dataset, device, device_num=1
):
    # study = optuna.create_study(direction="maximize")
    study = optuna.load_study(study_name=study_name, storage=storage)
    with open(stdout_file, "w") as f_out, open(stderr_file, "w") as f_err:
        with contextlib.redirect_stdout(f_out), contextlib.redirect_stderr(f_err):
            study.optimize(
                lambda trial: objective(trial, model, dataset, device),
                n_trials=300 // device_num,
            )
            best_trial = study.best_trial
            print(
                f"Best trial on {model} {dataset}: {best_trial.value} with {best_trial.params}"
            )


if __name__ == "__main__":
    DEVICES_NUM = 8
    DATASETS = ["ASSIST2009", "JUNYI", "NIPS2020"]
    # MODEL = ["DINA", "IRT", "KaNCD", "MCD", "MIRT", "NCDM"]
    MODEL = "MCD"
    # MODEL = ["MCD"]
    devices = ["cuda:{}".format(i) for i in range(DEVICES_NUM)]
    devices_index = 0
    tasks = []
    for dataset in DATASETS:
        optuna.create_study(
            study_name=f"{MODEL}-debias-{dataset}",
            storage=f"sqlite:///optuna/{MODEL}-debias-{dataset}.db",
            direction="maximize",
        )
    for i in range(DEVICES_NUM):
        for dataset in DATASETS:
            tasks.append(
                (
                    f"{MODEL}-debias-{dataset}",
                    f"sqlite:///optuna/{MODEL}-debias-{dataset}.db",
                    search_output_dir / f"{MODEL}_{dataset}_{i}.out",
                    search_output_dir / f"{MODEL}_{dataset}_{i}.err",
                    MODEL,
                    dataset,
                    devices[i],
                    DEVICES_NUM,
                )
            )

    with Pool(DEVICES_NUM * len(DATASETS) * 2) as pool:
        pool.starmap(run_wrapper, tasks)
        pool.close()
        pool.join()

    print("All Search finished!")
